#!/usr/bin/python
##    Copyright (c) 2006 Greg Robinson. Email: gregarobi@dev.java.net
##
##    Permission is hereby granted, free of charge, to any person obtaining a copy of this 
##    software and associated documentation files (the "Software"), to deal in the Software 
##    without restriction, including without limitation the rights to use, copy, modify,merge, 
##    publish, distribute, sublicense, and/or sell copies of the Software, and to permit 
##    persons to whom the Software is furnished to do so, subject to the following conditions: 
## 
##    The above copyright notice and this permission notice shall be included in all copies 
##    or substantial portions of the Software. 
## 
##    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, 
##    INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A 
##    PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT 
##    HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF 
##    CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR 
##    THE USE OR OTHER DEALINGS IN THE SOFTWARE. 
 
import math, numpy, random, time 
 
 
#=================================== 
##    Swarm creates swarm of "nParticles" particles in "nDimensions"-dimensional hyperspace 
##    with dimension lengths "spaceSize". Each particle has a position vector contained in 
##    "self.position" and a corresponding velocity vector contained in "self.velocity". 
##     
##    You can think of this as a swarm of bees, where each bee is adjusting its velocity 
##    towards the "best" bee in the swarm. The "best" bee is the one that, for example, 
##    found the most honey. 
## 
##    This optimization can be applied to any smooth and continuous function of any number 
##    of variables. "Swarm.fitnessFunc()" should be replaced with the applicable function 
##    by way of inheritence. 
## 
##    TODO:
##        -Convergence instead of absolute error.
##            Sometimes, and often, the optimum value is not known. Add functionality for
##            to find relative convergence of the swarm.
##        -Eliminate hard-coded maximum fitness currently on updateFitness. There are times
##            when less is better or when one is trying to find a specific value. 
##        -Neighborhood swarming. "self.metric" is currently unused, but will be used to find 
##            neighbor particles. This functionality is currently in progress. It should solve 
##            most of the local optima difficulties this currently has. 
##        -More flexible algorithm for updateVelocity(). 
##            Maybe add more weight options such as "crazyness"? 
##        -Faster method for updateVelocity natively using numpy? 
##            According to unittest, this is one of the slowest parts. 
##        -Comments! I got carried away with coding, so it might be difficult for 
##            other people to understand it. 
class Swarm: 
    def __init__(self, nDimensions, spaceSize, nParticles, learn_rate): 
        self.LEARN_RATE = learn_rate 
        self.nDimensions = nDimensions 
        self.spaceSize = spaceSize 
        self.nParticles = nParticles 
        self.position = numpy.random.random((nParticles,nDimensions)) * spaceSize 
        self.MAX_POSITION = numpy.ones(nDimensions) * spaceSize 
        self.MIN_POSITION = numpy.zeros(nDimensions) 
        self.velocity = numpy.random.uniform(-spaceSize, spaceSize, (nParticles,nDimensions)) 
        self.updateFitness() 
        self.updateMetric() 
 
    # This gets overwritten when child classes define their own fitness function 
    def fitnessFunc(self, particleIndex): 
        p = self.position[particleIndex] 
        x = p[0] 
        y = p[1] 
        return 0.5-(math.sin(math.sqrt(x**2 +y**2))**2-0.5)/ \
               (1.0+0.001*(x**2+y**2))**2
 
    def updateFitness(self): 
        self.fitness = numpy.zeros(self.nParticles, 'd') 
        for i in xrange(self.nParticles): 
            self.fitness[i] = self.fitnessFunc(i) 
        self.sorted = numpy.argsort(self.fitness).tolist() 
        self.sorted.reverse() 
        return self.fitnessFunc(self.sorted[0]) 
         
    def updatePosition(self): 
        position = self.position+self.velocity 
        position = numpy.minimum(position, self.MAX_POSITION) 
        position = numpy.maximum(position,self.MIN_POSITION) 
        self.position = position 
 
    def updateVelocity(self): 
	best = self.position[self.sorted[0]] 
        for dimension in xrange(self.nDimensions): 
            for pindex in xrange(self.nParticles): 
                self.velocity[pindex][dimension] = \
                            random.uniform(0, self.velocity[pindex][dimension]) + \
                            self.LEARN_RATE*random.uniform(0, \
                            best[dimension]-self.position[pindex][dimension]) 
 
    def updateMetric(self):
        self.position = numpy.array(self.position)
        n,m = self.position.shape
        delta = numpy.zeros((n,n),'d')
        for d in xrange(m):
            data = self.position[:,d]
            delta += (data - data[:,numpy.newaxis])**2
        self.metric = numpy.sqrt(delta)

    def getNeighbors(self,):
        pass

    def plotPoints(self, gp):
        #####
        #This function is hacked to pieces. Clean up.
        #####
        import Gnuplot
        corners = [[0,0,0],[0,self.spaceSize,0],[self.spaceSize,0,1]]
        if gp==0:
            gp = Gnuplot.Gnuplot(persist = 0)
            gp('set data style points')
            corners = Gnuplot.PlotItems.Data(corners, title=None)  # No title
            gp.plot(corners)
        else:
            data = [[self.position[i][0], self.position[i][1], self.fitnessFunc(i)] for i in xrange(self.nParticles)]
            plot = Gnuplot.PlotItems.Data(data, title=None)  # No title
            #gp.splot(plot,corners)

            ###########DEMO CODE
            gp.splot(corners,'sin(sqrt((x-pi)**2+(y-pi)**2)) / sqrt((x-pi)**2+(y-pi)**2)',plot)

        return gp 
 
    def printData(self): 
        #this is just development/debug code 
        print 'position\n', self.position 
        print 
        print 'velocity\n', self.velocity 
        print 
        print 'metric\n', self.metric 
        print 
        print 'fitness\n', self.fitness 
        print 
        print 'sort\n', self.sorted 
        print 
        print 'best\n', self.fitnessFunc(self.sorted[0]) 
        print self.position[self.sorted[0]] 
#====================================
 
def main(): 
    s = Swarm(nDimensions=2,spaceSize=100,nParticles=12,learn_rate=.9) 
    MAX_ERROR = 0.00001
    errorFunc = lambda :2 - s.updateFitness()
    error = errorFunc() 
    print error
    gp = s.plotPoints(0)
    while error>=MAX_ERROR: 
        s.updateVelocity() 
        s.updatePosition() 
        gp = s.plotPoints(gp)
        error = errorFunc() 
        print error
    print error 
    s.printData() 
    time.sleep(2)
 
if __name__ == '__main__': 
    for count in xrange(10): 
        main()
